package com.bigdata.kafka.utils;

import com.bigdata.auth.SourceType;
import org.apache.commons.lang3.StringUtils;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;

/**
 * @author xin.ding
 * @date 2021/4/6 14:04
 */
public class KafkaAuth {

    public static void buildKafkaSaslEnv(Properties properties, String username, String password) {
        properties.put("security.protocol", "SASL_PLAINTEXT");
        properties.put("sasl.mechanism", "PLAIN");
        properties.put("sasl.jaas.config",
                String.format("org.apache.kafka.common.security.plain.PlainLoginModule required username=\"%s\" password=\"%s\";"
                        , username
                        , password));
    }

    /**
     * @param properties 配置文件
     * @param krb5Conf   krb5.conf配置文件 D:\\tmp\\krb5.conf
     * @param jaasConf   JAAS配置文件 D:\\tmp\kafka-jaas.conf
     */
    public static void buildKerberosEnv(Properties properties, String krb5Conf, String jaasConf) {
        System.setProperty("java.security.krb5.conf", krb5Conf);
        System.setProperty("java.security.auth.login.config", jaasConf);
        properties.put("security.protocol", "SASL_PLAINTEXT");
        properties.put("sasl.kerberos.service.name", "kafka");
        //Kerberos服务名,对应于kafka-jaas.conf中的serviceName
//        properties.put("sasl.kerberos.service.name", "kafka");
        //服务器端的安全协议,此处是Kerberos认证,所以是SASL_PLAINTEXT
//        properties.put("security.protocol", "SASL_PLAINTEXT");
    }

    /**
     * 生成jass文件
     *
     * @param type
     * @param principal
     * @param keytab
     * @return
     */
    public static String buildJassFile(SourceType type, String principal, String keytab) {
        File file = newJassFile();

        try {
            List<String> lines = new ArrayList<>(10);

            if (type == SourceType.KAFKA) {
                lines.add("KafkaClient {");
            } else if (type == SourceType.HBASE) {
                lines.add("HbaseClient {");
            } else if (type == SourceType.HIVE) {
                lines.add("com.sun.security.jgss.krb5.initiate {");
            }

            lines.add(" com.sun.security.auth.module.Krb5LoginModule required");
            if (type == SourceType.HBASE) {
                lines.add(" serviceName=\"kafka\"");
            }
            lines.add(" useTicketCache=true");
            lines.add(" useKeyTab=true");
            lines.add(" principal=\"" + principal + "\"");
            lines.add(" keyTab=\"" + keytab + "\";");
            lines.add("};");

            lines.add("Client {");
            lines.add(" com.sun.security.auth.module.Krb5LoginModule required");
            lines.add(" useTicketCache=true");
            lines.add(" useKeyTab=true");
            lines.add(" principal=\"" + principal + "\"");
            lines.add(" keyTab=\"" + keytab + "\";");
            lines.add("};");

            FileWriter fw = new FileWriter(file.getAbsoluteFile(), true);
            fw.append(StringUtils.join(lines, "\n"));
            fw.close();

        } catch (IOException e) {
            e.printStackTrace();
        }


        return file.getAbsolutePath();
    }

    /**
     * 创建临时jass文件
     *
     * @return File
     */
    private static File newJassFile() {
        final File file;

        try {
            Path path = Files.createTempFile("jaas-", ".conf");
            file = path.toFile();
            file.deleteOnExit();
        } catch (IOException e) {
            throw new RuntimeException("unable to generate a JAAS configuration file", e);
        }

        return file;
    }
}
